#include "FEM2DVector.h"
#include <time.h>
#include <fstream>
using namespace std;

CFEM2DVector::CFEM2DVector(int n):
	m_nSizeVec(n)
{
	cout << "=============================================================" << endl;
	cout << "****                                                     ****" << endl;
	cout << "****                   GenFEM Version 0.1                ****" << endl;
	cout << "****                                                     ****" << endl;
	cout << "=============================================================" << endl;
	cout << "> About: \tGenFEM is a general package for finite element\n"
		<< "         \twritten in C++" << endl;
	cout << "> Author: \tYizhao Wan, Qingdao Institute of Marine Geology\n"
		<< "           \tChina Geological Survey, Shandong, China " << endl;
	cout << "> Release: \ton Jan 11, 2021" << endl;
	cout << "> Website: " << endl;
	cout << "> Email: \twanyizhao@mail.cgs.gov.cn" << endl;
	cout << "=============================================================" << endl;
	
}

CFEM2DVector::CFEM2DVector()
{

}
CFEM2DVector::~CFEM2DVector()
{

}
void CFEM2DVector::Init()
{
	
	cout << "GenFEM Initializing..." << endl;
	for (unsigned int i = 0; i < m_Mesh.getAllBoundEdges().size(); ++i)
	{
		for (unsigned int j = 0; j < m_Mesh.getAllBoundEdges()[i].size(); ++j)
		{
			m_Mesh.getAllBoundEdges()[i][j].value.resize(m_nSizeVec);
		}
	}
	for (unsigned int i = 0; i < m_Mesh.getAllBoundNodes().size(); ++i)
	{
		for (unsigned int j = 0; j < m_Mesh.getAllBoundNodes()[i].size(); ++j)
		{
			m_Mesh.getAllBoundNodes()[i][j].value.resize(m_nSizeVec);
		}
	}

}

void CFEM2DVector::assignBoundEdgesValue(int bound_index, int vec_index, BOUNDARY_TYPE type, double(*value)(CPoint, double t), double(*r)(CPoint, double t))
{
	cout << "Assign boundary edges value for boundary\t" << bound_index + 1 << endl;
	if (vec_index >= m_nSizeVec)
	{
		cout<<"Error: boundary's dimension is greater than the maximum dimension of the unknowns"<<endl;
		exit(100);
		
	}
	if (bound_index>=m_Mesh.getAllBoundEdges().size())
	{
		cout << "Error: boundary's index is greater than the maximum number of boundary" << endl;
		exit(100);
	}
	m_Mesh.assignBoundEdgesValue(m_Mesh.getAllBoundEdges()[bound_index], vec_index, type, value, r);
	
}

void CFEM2DVector::assignBoundNodesValue(int bound_index, int vec_index, BOUNDARY_TYPE type, double(*value)(CPoint, double t))
{
	cout << "Assign boundary nodes value for " << vec_index + 1 << " value of boundary " << bound_index + 1 << endl;
	if (vec_index >= m_nSizeVec)
	{
		cout << "Error: boundary's dimension is greater than the maximum dimension of the unknowns" << endl;
		exit(100);
		
	}
	if (bound_index >= m_Mesh.getAllBoundNodes().size())
	{
		cout << "Error: boundary's index is greater than the maximum number of boundary" << endl;
		exit(100);
	}
	m_Mesh.assignBoundNodesValue(m_Mesh.getAllBoundNodes()[bound_index], vec_index, type, value);
	

}

void CFEM2DVector::setNumVec(int n)
{
	m_nSizeVec = n;
}

int CFEM2DVector::getNumVec()
{
	return m_nSizeVec;
}

void CFEM2DVector::treatBoundaryConditions(double(*coe)(CPoint, double t), double t, SpMat & A, Rsv & b)
{
	clock_t start, end;
	cout << "Treat boundary conditions" << endl;
	start = clock();
	for (unsigned int k = 0; k < m_Mesh.getAllBoundNodes().size(); ++k)
	{
		for (unsigned int l = 0; l < m_Mesh.getAllBoundNodes()[k].size(); ++l)
		{
			if (Dirichlet == m_Mesh.getAllBoundNodes()[k][l].type)
			{
				//set the A(i,i)=1 and A(i,j)=0, b=g(Pb(i)) to handle Dirichlet boundary
				//It is needed a loop over the column, and the manipulation of matrix is severely time consuming 
				//So we use the the second way to handle Dirichlet boundary

				int i = m_Mesh.getAllBoundNodes()[k][l].glo_node_index;
				for (unsigned int num = 0; num < m_Mesh.getAllBoundNodes()[k][l].mark.size(); ++num)
				{
					int index = m_Mesh.getAllBoundNodes()[k][l].mark[num];
					for (int j = 0; j < A.cols(); ++j)
					{
						A.coeffRef(i + index * m_Mesh.getPbTrial().size(), j) = 0;
					}
					A.coeffRef(i + index * m_Mesh.getPbTrial().size(), i + index * m_Mesh.getPbTrial().size()) = 1;
					b[i + index * m_Mesh.getPbTrial().size()] = m_Mesh.getAllBoundNodes()[k][l].value[index](m_Mesh.getPbTrial()[i], t);
				}

				//another way to handle the Dirichlet boundary
				//b(i)=b=g(Pb(i))*A(i,i)*large_number
				//A(i,i)=A(i,i)*large_number
				//It is only needed to manipulate the diagonal element of the matrix

				/*int i = m_Mesh.getAllBoundNodes()[k][l].glo_node_index;
				double large_number = 1e50;
				for (unsigned int num=0;num< m_Mesh.getAllBoundNodes()[k][l].mark.size();++num)
				{
					int index = m_Mesh.getAllBoundNodes()[k][l].mark[num];
					//Eigen::SparseMatrix<double>::InnerIterator it(A, i+num* m_Mesh.getPbTrial().size());
					double A_ii = A.coeffRef(i + index * m_Mesh.getPbTrial().size(), i + index * m_Mesh.getPbTrial().size());
					b(i+ index * m_Mesh.getPbTrial().size()) = A_ii *large_number*m_Mesh.getAllBoundNodes()[k][l].value[index](m_Mesh.getPbTrial()[i], t);
					A.coeffRef(i+ index * m_Mesh.getPbTrial().size(), i+ index * m_Mesh.getPbTrial().size()) = A_ii * large_number;
				}*/
				
				
			}
		}

	}

	for (unsigned int k = 0; k < m_Mesh.getAllBoundEdges().size(); ++k)
	{
		for (unsigned int l = 0; l < m_Mesh.getAllBoundEdges()[k].size(); ++l)
		{
			if (Neumann == m_Mesh.getAllBoundEdges()[k][l].type)
			{
				int nk = m_Mesh.getAllBoundEdges()[k][l].ele_index;
				int p1 = m_Mesh.getAllBoundEdges()[k][l].p.first;
				int p2 = m_Mesh.getAllBoundEdges()[k][l].p.second;

				//This vertices is the boundary edge's two nodes
				//It is for the generation of Gauss points on the edges
				vector<CPoint> vertices_ede{
					m_Mesh.getPbTest()[p1],
					m_Mesh.getPbTest()[p2]
				};
				vector<CPoint>Gauss_nodes;
				vector<double>Gauss_wights;
				generateGaussLocal1D(vertices_ede, m_GaussType, Gauss_nodes, Gauss_wights);

				//The basis function needs the vertices of the element not the edge
				//So it is needed to get the nk's nodes
				vector<CPoint> vertices;//vertices (coordinates) of nth element(mesh nodes not finite element nodes)		
				for (int i = 0; i < m_Mesh.getElements()[nk].get_ele_node_num(); ++i)
				{
					vertices.push_back(m_Mesh.getPoints()[m_Mesh.getElements()[nk].getIds()[i]]);
				}
				for (unsigned int num = 0; num < m_Mesh.getAllBoundEdges()[k][l].mark.size(); ++num)
				{
					int index = m_Mesh.getAllBoundEdges()[k][l].mark[num];
					for (unsigned int beta = 0; beta < m_Mesh.getTbTest()[nk].getIds().size(); ++beta)
					{
						double ini_value = GaussQuad1DBoundTrial(coe, t, m_Mesh.getAllBoundEdges()[k][l].value[index], Gauss_wights, Gauss_nodes, vertices, m_Mesh.getTbTrial()[nk].getBasisType(ELE_BASIS), beta, 0, 0);
						b(m_Mesh.getTbTest()[nk].getIds()[beta]+ index * m_Mesh.getPbTrial().size()) += ini_value;
					}
				}
			}
			else if (Robin == m_Mesh.getAllBoundEdges()[k][l].type)
			{
				int nk = m_Mesh.getAllBoundEdges()[k][l].ele_index;
				int p1 = m_Mesh.getAllBoundEdges()[k][l].p.first;
				int p2 = m_Mesh.getAllBoundEdges()[k][l].p.second;

				//This vertices is the boundary edge's two nodes
				//It is for the generation of Gauss points on the edges
				vector<CPoint> vertices_ede{
					m_Mesh.getPbTest()[p1],
					m_Mesh.getPbTest()[p2]
				};
				vector<CPoint>Gauss_nodes;
				vector<double>Gauss_wights;
				generateGaussLocal1D(vertices_ede, m_GaussType, Gauss_nodes, Gauss_wights);

				//The basis function needs the vertices of the element not the edge
				//So it is needed to get the nk's nodes
				vector<CPoint> vertices;//vertices (coordinates) of nth element(mesh nodes not finite element nodes)		
				for (int i = 0; i < m_Mesh.getElements()[nk].get_ele_node_num(); ++i)
				{
					vertices.push_back(m_Mesh.getPoints()[m_Mesh.getElements()[nk].getIds()[i]]);
				}
				for (unsigned int num = 0; num < m_Mesh.getAllBoundEdges()[k][l].mark.size(); ++num)
				{
					int index = m_Mesh.getAllBoundEdges()[k][l].mark[num];
					for (unsigned int beta = 0; beta < m_Mesh.getTbTest()[nk].getIds().size(); ++beta)
					{
						double ini_value = GaussQuad1DBoundTrial(coe, t, m_Mesh.getAllBoundEdges()[k][l].value[index], Gauss_wights, Gauss_nodes, vertices, m_Mesh.getTbTrial()[nk].getBasisType(ELE_BASIS), beta, 0, 0);
						b(m_Mesh.getTbTest()[nk].getIds()[beta]+ index * m_Mesh.getPbTrial().size()) += ini_value;

						for (unsigned int alpha = 0; alpha < m_Mesh.getTbTrial()[nk].getIds().size(); ++alpha)
						{
							double ini_value = GaussQuad1DBoundTrialTest(coe, t, m_Mesh.getAllBoundEdges()[k][l].r[index], Gauss_wights, Gauss_nodes, vertices,
								m_Mesh.getTbTrial()[nk].getBasisType(ELE_BASIS), alpha, 0, 0, m_Mesh.getTbTest()[nk].getBasisType(ELE_BASIS), beta, 0, 0);
							int i = m_Mesh.getTbTest()[nk].getIds()[beta];
							int j = m_Mesh.getTbTrial()[nk].getIds()[alpha];
							A.coeffRef(i+ index * m_Mesh.getPbTrial().size(), j+ index * m_Mesh.getPbTrial().size()) += ini_value;

						}
					}
				}
			}
		}

	}
	end = clock();
	cout << "Time consumption of boundary conditions treatment :\t" << (double)(end - start) / CLOCKS_PER_SEC << " sec" << endl << endl;

}
SpMat CFEM2DVector::matrix_hstack(SpMat m1, SpMat m2)
{
	/*SpMat m1_hstack_m2(m1.rows(), m1.cols() + m2.cols());
	m1_hstack_m2.leftCols(m1.cols()) = m1;
	m1_hstack_m2.rightCols(m2.cols()) = m2;
	return m1_hstack_m2;*/

	SparseMatrix<double> M(m1.rows(), m1.cols()+m2.cols());
	M.reserve(m1.nonZeros() + m2.nonZeros());
	for (Index c = 0; c < m1.cols(); ++c)
	{
		M.startVec(c); // Important: Must be called once for each rows before inserting!
		for (SparseMatrix<double>::InnerIterator itL(m1, c); itL; ++itL)
			M.insertBack(itL.row(), c) = itL.value();		
	}
	for (Index c = 0; c <m2.cols(); ++c)
	{
		M.startVec(c+m1.cols()); // Important: Must be called once for each rows before inserting!		
		for (SparseMatrix<double>::InnerIterator itC(m2, c); itC; ++itC)
			M.insertBack(itC.row(), c + m1.cols()) = itC.value();
	}
	M.finalize();
	return M;


}
Eigen::SparseMatrix<double> CFEM2DVector::matrix_vstack(SpMat m1, SpMat m2)
{
	/*Eigen::SparseMatrix<double, Eigen::RowMajor> m1_vstack_m2(m1.rows() + m2.rows(), m1.cols());
	m1_vstack_m2.topRows(m1.rows()) = m1;
	m1_vstack_m2.bottomRows(m2.rows()) = m2;
	cout << m1_vstack_m2.IsRowMajor << endl;
	return m1_vstack_m2;*/

	SparseMatrix<double> M(m1.rows() + m2.rows(), m1.cols());
	M.reserve(m1.nonZeros() + m2.nonZeros());
	for (Index c = 0; c < m1.cols(); ++c)
	{
		M.startVec(c); // Important: Must be called once for each column before inserting!
		for (SparseMatrix<double>::InnerIterator itL(m1, c); itL; ++itL)
			M.insertBack(itL.row(), c) = itL.value();
		for (SparseMatrix<double>::InnerIterator itC(m2, c); itC; ++itC)
			M.insertBack(itC.row() + m1.rows(), c) = itC.value();
	}
	M.finalize();
	return M;
}

Rsv CFEM2DVector::vector_stack(Rsv b1, Rsv b2)
{
	Rsv b12(b1.rows() + b2.rows());
	b12 << b1, b2;
	return b12;
}
void CFEM2DVector::Export2VTK(string filename, Rsv x)
{
	cout << "\nExport results to vtk visualized with Paraview" << endl;
	map<ELE_TYPE, int> type_ele_vtk = {
		{ELE_3,5},
		{ELE_4,9},
		{ELE_6,22},
		{ELE_9,28}
	};
	ofstream f(filename);
	f << "# vtk DataFile Version 3.0" << endl;
	f << "Result" << endl;
	f << "ASCII" << endl;
	f << "DATASET UNSTRUCTURED_GRID" << endl;

	f << "POINTS \t" << m_Mesh.getPbTest().size() << "\tdouble" << endl;
	for (unsigned int i = 0; i < m_Mesh.getPbTest().size(); ++i)
	{
		f << m_Mesh.getPbTest()[i].getX() << "\t" << m_Mesh.getPbTest()[i].getY() << "\t" << m_Mesh.getPbTest()[i].getZ() << endl;
	}
	f << "CELLS \t " << m_Mesh.getTbTest().size() << "\t" << (m_Mesh.getTbTest()[0].getIds().size() + 1)* m_Mesh.getTbTest().size() << endl;
	for (unsigned int i = 0; i < m_Mesh.getTbTest().size(); ++i)
	{

		f << m_Mesh.getTbTest()[i].getIds().size() << "\t";

		if (m_Mesh.getTbTest()[i].get_ele_type() == ELE_3 || m_Mesh.getTbTest()[i].get_ele_type() == ELE_4)
		{
			for (unsigned int j = 0; j < m_Mesh.getTbTest()[i].getIds().size(); ++j)
			{
				f << m_Mesh.getTbTest()[i].getIds()[j] << "\t";
			}
		}
		else if (m_Mesh.getTbTest()[i].get_ele_type() == ELE_6)
		{
			for (unsigned int j = 0; j < m_Mesh.getTbTest()[i].getIds().size(); j += 2)
			{
				f << m_Mesh.getTbTest()[i].getIds()[j] << "\t";
			}
			for (unsigned int j = 1; j < m_Mesh.getTbTest()[i].getIds().size(); j += 2)
			{
				f << m_Mesh.getTbTest()[i].getIds()[j] << "\t";
			}
		}
		else if (m_Mesh.getTbTest()[i].get_ele_type() == ELE_9)
		{
			for (unsigned int j = 0; j < m_Mesh.getTbTest()[i].getIds().size() - 1; j += 2)
			{
				f << m_Mesh.getTbTest()[i].getIds()[j] << "\t";
			}
			for (unsigned int j = 1; j < m_Mesh.getTbTest()[i].getIds().size() - 1; j += 2)
			{
				f << m_Mesh.getTbTest()[i].getIds()[j] << "\t";
			}
			f << m_Mesh.getTbTest()[i].getIds()[8] << "\t";
		}
		f << endl;

	}
	f << "CELL_TYPES\t" << m_Mesh.getTbTest().size() << endl;

	for (unsigned int i = 0; i < m_Mesh.getTbTest().size(); ++i)
	{
		f << type_ele_vtk[m_Mesh.getTbTest()[i].get_ele_type()] << endl;
	}

	f << "POINT_DATA\t" << m_Mesh.getPbTest().size() << endl;

	f << "VECTORS s_fem double" << endl;	
	for (unsigned int i = 0; i < m_Mesh.getPbTest().size(); ++i)
	{

		f << x(i) << "\t";
		f << x(i + m_Mesh.getPbTest().size()) << "\t";
		f << 0 << endl;

	}
	f.close();

}
